# Evaluate of SIFA
import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import numpy as np
import matplotlib

matplotlib.use('Agg')
from torch.utils.data import DataLoader
import torch
from model import SIFA
from dataset.dataset import SingleDataset
import cv2
from dataset.dataset import parse_config
from tqdm import tqdm
from medpy import metric
from myMetrics import Meterics
config = "/home/liukai/projects/SIFA/config/train4label.cfg"
config = parse_config(config)
exp_name = config['train']['exp_name']

def dice_eval(predict,label,num_classes):
    '''
    Dice=2*交集/并集
    如果target中没有出现该类别，预测结果也没有该类别，那么该类别的dice=100%
    '''
    #Computer Dice coefficient
    dice = np.zeros(num_classes)
    eps = 1e-7#防止分母为0
    for c in range(num_classes):
        pre_sum=(predict==c)
        true_sum=(label==c)
        inter = 2.0 * (np.sum((predict==c)*(label==c),dtype=np.float16))
        p_sum = np.sum(predict==c,dtype=np.float32)
        gt_sum = np.sum(label==c,dtype=np.float32)
        dice[c] = (inter+eps)/(p_sum+gt_sum+eps)
    return dice[1:]

def assd_eval(predict, label, num_classes):
    # Average Symmetric Surface Distance (ASSD)
    assd_all = np.zeros(num_classes)
    for c in range(num_classes):
        reference = (label == c) * 1
        result = (predict == c) * 1
        assd_all[c] = metric.binary.assd(result, reference)
    return assd_all[1:]


def create_visual_anno(anno):
    assert np.max(anno) < 7  # only 7 classes are supported, add new color in label2color_dict
    label2color_dict = {  # BGR
        0: [0, 0, 0],
        1: [0, 0, 255],  # 红色
        2: [0, 255, 0],  # 绿色
        3: [122, 209, 255],  # 黄色
        4: [255, 215, 0],  # 浅蓝色
        5: [160, 32, 100],
        6: [255, 64, 64],
        7: [139, 69, 19],
    }
    # visualize
    visual_anno = np.zeros((anno.shape[0], anno.shape[1], 3), dtype=np.uint8)
    for i in range(visual_anno.shape[0]):  # i for h
        for j in range(visual_anno.shape[1]):
            color = label2color_dict[anno[i, j]]
            visual_anno[i, j, 0] = color[0]
            visual_anno[i, j, 1] = color[1]
            visual_anno[i, j, 2] = color[2]

    return visual_anno


def norm_01(image):
    mn = np.min(image)
    mx = np.max(image)
    image = (image - mn) / (mx - mn).astype(np.float32)
    return image


def save_img(image):
    image = norm_01(image)
    image = (image * 255).astype(np.uint8)
    return image


device = torch.device('cuda:{}'.format(config['test']['gpu']))
test_path = config['test']['test_path']
num_classes = config['test']['num_classes']
sifa_model = SIFA(config).to(device)
sifa_model.load_state_dict(torch.load('{}'.format(config['test']['test_model']),
                                      map_location=torch.device('cuda:{}'.format(config['test']['gpu']))))
sifa_model.eval()
# test dataset
test_dataset = SingleDataset(test_path)
batch_size = config['test']['batch_size']
test_loader = DataLoader(test_dataset, batch_size, shuffle=False)
myMeteric=Meterics(config['test']['num_classes'])

# test
all_batch_dice = []
all_batch_assd = []
with torch.no_grad():
    for it, (xt, xt_label) in tqdm(enumerate(test_loader)):
        xt = xt.to(device)
        xt_label = xt_label.numpy().squeeze().astype(np.uint8)

        output = sifa_model.test_seg(xt).detach()
        output = output.squeeze(0)
        output = torch.argmax(output, dim=0)
        output = output.cpu().numpy()

        xt = xt.detach().cpu().numpy().squeeze()
        gt = xt_label.reshape(256, 256).astype(np.uint8)
        output = output.squeeze()

        xt = save_img(xt)

        output_vis = create_visual_anno(output)
        gt_vis = create_visual_anno(gt)
        results = config['test']['predict_image_path']
        if (not os.path.exists(results)):
            os.makedirs(results)
            print(results)
        cv2.imwrite('{}/xt-{}.jpg'.format(results, it + 1), xt)
        cv2.imwrite('{}/gt-{}.jpg'.format(results, it + 1), gt_vis)
        cv2.imwrite('{}/output-{}.jpg'.format(results, it + 1), output_vis)

        myMeteric.add_batch(xt_label,output)
        try:
            one_case_assd = assd_eval(output, xt_label, num_classes)
        except:
            continue


'''
    脾（spleen）：1         红色
    右肾（right kidney）：2  绿色
    左肾（left kidney）：3   黄色
    肝：4                   浅蓝色
'''

all_batch_assd = np.array(all_batch_assd)
mean_assd = np.mean(all_batch_assd, axis=0)
mean_dice=myMeteric.get_dice()
mean_precision=myMeteric.get_precision()
mean_recall=myMeteric.get_recall()
# 输出

print(all_batch_dice)
std_assd = np.std(all_batch_assd, axis=0)
print('-----------')
print('MYO||LV||RV')
print('Dice mean:{}'.format(mean_dice))
print('precision mean:{}'.format(mean_precision))
print('recall mean:{}'.format(mean_recall))
print('ASSD mean:{}'.format(mean_assd))
print('ASSD std:{}'.format(std_assd))
print('total mean assd:', np.mean(mean_assd))
print('-----------')

'''

'''

